#include <mpi.h>
#include <windows.h>
#include <stdio.h>
#include <string.h>

#define BREAK_HERE() __asm __emit 0xF1
#define BUFFER_SIZE 1024

static SOCKET fd;

static int Send (SOCKET fd, char *buffer, int len)
{
    return send (fd, buffer, len, 0);
}

static int Receive (SOCKET fd, char *buffer, int len)
{
    return recv (fd, buffer, len, 0);
}

static SOCKET Connect (const char *hostname, int port)
{
    struct sockaddr_in    server_addr_in;
    struct sockaddr       *server_addr = NULL;
    struct hostent        *server = NULL;
    SOCKET                fd;
    int                   length;

    /*
     * Initialize Winsock
     */
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2,2), &wsaData) != NO_ERROR)
    {
        printf ("Error Starting Winsock\n");
        exit (1);
    }

    /*
     * Open the Socket endpoint.
     */
    if ((fd = socket (AF_INET, SOCK_STREAM, 0)) == INVALID_SOCKET)
    {
        return INVALID_SOCKET;
    }

    memset ((char *)&server_addr_in, '\0', sizeof (server_addr_in));
    server = gethostbyname (hostname);    

    if ((char *) server == NULL)
    {
        printf ("Invalid Server name %s\n", hostname);
        return INVALID_SOCKET;
    }

    server_addr_in.sin_family = AF_INET;
    memmove ((char *)&server_addr_in.sin_addr.s_addr, (char *)server->h_addr,
                server->h_length);
    server_addr_in.sin_port = htons (port);
    server_addr = (struct sockaddr *)&server_addr_in;
    length = sizeof (server_addr_in);

    /*
     * Connect to the server.
     */
    if (connect (fd, server_addr, length) < 0)
    {
        printf ("Client unable to connect to server\n");
        return INVALID_SOCKET;
    }

    return fd;
}

static void Disconnect (SOCKET fd)
{
    closesocket (fd);
    WSACleanup();
}

static int Break(char *mpiCallName)
{
    static char buffer[50];
    static int connected = 1;
    int length;
    
    if (!connected)
        return 0;

    // Append a new line
    strcpy_s(buffer, sizeof(buffer), mpiCallName);
    strcat_s(buffer, sizeof(buffer), "\n");

    // Do a send and a receive
    Send(fd, buffer, strlen(buffer));
    length = Receive(fd, buffer, 1024);
    buffer[length] = '\0';

    // Return the result.
    // 1 = Break
    // 0 = Don't break
    // 2 = Break once, and continue in the future
    if (!strcmp(buffer, "1\n"))
    {
        return 1;
    }
    else if (!strcmp(buffer, "0\n"))
    {
        return 0;
    }
    else
    {
        connected = 0;
        return 1;
    }
}

int MPIAPI MPI_Init(int *argc, char ***argv)
{
    int rank;
    char rankStr[100];
    char pidStr[100];
    char *hostname;
    size_t hostnameLength;
    int port;
    char *portStr;
    size_t portLength;
    int retval;

    // Read the server info from the environment.
    _dupenv_s(&hostname, &hostnameLength, "ISP_HOSTNAME");
    _dupenv_s(&portStr, &portLength, "ISP_PORT");
    port = atoi(portStr);

    // Connect to the server.
    fd = Connect(hostname, port);
    if (fd == INVALID_SOCKET)
    {
        exit(1);
    }

    // Issue the MPI commands.
    retval = PMPI_Init(argc, argv);
    PMPI_Comm_rank(MPI_COMM_WORLD, &rank);

    // First send the rank.
    sprintf_s(rankStr, sizeof(rankStr), "%d\n", rank);
    Send(fd, rankStr, strlen(rankStr));

    // Next send the PID.
    sprintf_s(pidStr, sizeof(pidStr), "%d\n", GetCurrentProcessId());
    Send(fd, pidStr, strlen(pidStr));

    return retval;
}

int MPIAPI MPI_Send (void * buffer, int count, MPI_Datatype datatype, int dest, 
                        int tag, MPI_Comm comm)
{
    if (Break("MPI_Send")) BREAK_HERE();
    return PMPI_Send(buffer, count, datatype, dest, tag, comm);
}

int MPIAPI MPI_Rsend (void * buffer, int count, MPI_Datatype datatype, int dest, 
                        int tag, MPI_Comm comm)
{
    if (Break("MPI_Rsend")) BREAK_HERE();
    return PMPI_Rsend(buffer, count, datatype, dest, tag, comm);
}

int MPIAPI MPI_Isend (void *buf, int count, MPI_Datatype datatype, int dest,
        int tag, MPI_Comm comm, MPI_Request *request)
{
    if (Break("MPI_Isend")) BREAK_HERE();
    return PMPI_Isend(buf, count, datatype, dest, tag, comm, request);
}    

#if 0

int MPIAPI MPI_Send_init (void *buf, int count, MPI_Datatype datatype, int dest,
        int tag, MPI_Comm comm, MPI_Request *request)
{
    if (Break("MPI_Send_init")) BREAK_HERE();
    return PMPI_Send_init(buf, count, datatype, dest, tag, comm, request);
}    

#endif

int MPIAPI MPI_Irecv (void *buf, int count, MPI_Datatype datatype, int source,
        int tag, MPI_Comm comm, MPI_Request *request)
{
    static char buffer[8];
    int length;
    int src;

    if (Break("MPI_Irecv")) BREAK_HERE();
    length = Receive(fd, buffer, sizeof(buffer));
    buffer[length] = '\0';
    src = atoi(buffer);
    if (src < 0)
        src = source;
    return PMPI_Irecv(buf, count, datatype, src, tag, comm, request);
}

int MPIAPI MPI_Recv (void *buf, int count, MPI_Datatype datatype, int source,
                int tag, MPI_Comm comm, MPI_Status *status)
{
    static char buffer[8];
    int length;
    int src;

    if (Break("MPI_Recv")) BREAK_HERE();
    length = Receive(fd, buffer, sizeof(buffer));
    buffer[length] = '\0';
    src = atoi(buffer);
    if (src < 0)
        src = source;
    return PMPI_Recv(buf, count, datatype, src, tag, comm, status);
}

#if 0

int MPIAPI MPI_Recv_init (void *buf, int count, MPI_Datatype datatype, int src,
        int tag, MPI_Comm comm, MPI_Request *request)
{
    if (Break("MPI_Recv_init")) BREAK_HERE();
    return PMPI_Recv_init(buf, count, datatype, src, tag, comm, request);
}    

#endif

int MPIAPI MPI_Probe(int source, int tag, MPI_Comm comm, MPI_Status *status)
{
    static char buffer[8];
    int length;
    int src;

    if (Break("MPI_Probe")) BREAK_HERE();
    length = Receive(fd, buffer, sizeof(buffer));
    buffer[length] = '\0';
    src = atoi(buffer);
    if (src < 0)
        src = source;
    return PMPI_Probe(src, tag, comm, status);
}

int MPIAPI MPI_Iprobe(int source, int tag, MPI_Comm comm, int *flag, MPI_Status *status)
{
    static char buffer[8];
    int length;
    int src;

    if (Break("MPI_Iprobe")) BREAK_HERE();
    length = Receive(fd, buffer, sizeof(buffer));
    buffer[length] = '\0';
    src = atoi(buffer);
    if (src < 0)
        src = source;
    return PMPI_Iprobe(src, tag, comm, flag, status);
}

int MPIAPI MPI_Wait (MPI_Request *request, MPI_Status *status)
{
    if (Break("MPI_Wait")) BREAK_HERE();
    return PMPI_Wait(request, status);
}

int MPIAPI MPI_Test (MPI_Request *request, int *flag, MPI_Status *status)
{
    if (Break("MPI_Test")) BREAK_HERE();
    return PMPI_Test(request, flag, status);
}

int MPIAPI MPI_Request_free (MPI_Request *request)
{
    if (Break("MPI_Request_free")) BREAK_HERE();
    return PMPI_Request_free(request);
}

#if 0
int MPIAPI MPI_Bsend (void * buffer, int count, MPI_Datatype datatype, int dest, 
                        int tag, MPI_Comm comm)
{
    if (Break("MPI_Bsend")) BREAK_HERE();
    return PMPI_Bsend(buffer, count, datatype, dest, tag, comm);
}
#endif

int MPIAPI MPI_Ssend (void *buffer, int count, MPI_Datatype datatype, int dest, 
                        int tag, MPI_Comm comm)
{
    if (Break("MPI_Ssend")) BREAK_HERE();
    return PMPI_Ssend(buffer, count, datatype, dest, tag, comm);
}

int MPIAPI MPI_Sendrecv (void * sendbuf, int sendcount, MPI_Datatype sendtype, 
    int dest, int sendtag, void * recvbuf, int recvcount, 
    MPI_Datatype recvtype, int source, int recvtag, MPI_Comm comm, 
                        MPI_Status *status)
{
    if (Break("MPI_Sendrecv")) BREAK_HERE();
    return PMPI_Sendrecv(sendbuf, sendcount, sendtype, dest, sendtag, recvbuf,
               recvcount, recvtype, source, recvtag, comm, status);
}

int MPIAPI MPI_Start (MPI_Request *request)
{
    if (Break("MPI_Start")) BREAK_HERE();
    return PMPI_Start(request);
}

int MPIAPI MPI_Startall (int count, MPI_Request *array_of_requests)
{
    if (Break("MPI_Startall")) BREAK_HERE();
    return PMPI_Startall(count, array_of_requests);
}

int MPIAPI MPI_Barrier (MPI_Comm comm)
{
    if (Break("MPI_Barrier")) BREAK_HERE();
    return PMPI_Barrier(comm);
}
int MPIAPI MPI_Cart_create ( MPI_Comm comm_old, int ndims, int *dims, int *periods, 
                      int reorder, MPI_Comm *comm_cart )
{
    if (Break("MPI_Cart_create")) BREAK_HERE();
return PMPI_Cart_create(comm_old, ndims, dims, periods, reorder, comm_cart);
}

int MPIAPI MPI_Comm_create ( MPI_Comm comm, MPI_Group group, MPI_Comm *comm_out )
{
    if (Break("MPI_Comm_create")) BREAK_HERE();
    return PMPI_Comm_create(comm, group, comm_out);
}

int MPIAPI MPI_Comm_dup ( MPI_Comm comm, MPI_Comm *comm_out )
{
    if (Break("MPI_Comm_dup")) BREAK_HERE();
    return PMPI_Comm_dup(comm, comm_out);
}

int MPIAPI MPI_Comm_split ( MPI_Comm comm, int color, int key, MPI_Comm *comm_out )
{
    if (Break("MPI_Comm_split")) BREAK_HERE();
    return PMPI_Comm_split(comm, color, key, comm_out);
}

int MPIAPI MPI_Comm_free ( MPI_Comm *comm)
{
    if (Break("MPI_Comm_free")) BREAK_HERE();
    return PMPI_Comm_free(comm);
}

int MPIAPI MPI_Bcast (void *buffer, int count, MPI_Datatype datatype, int root,
        MPI_Comm comm)
{
    if (Break("MPI_Bcast")) BREAK_HERE();
    return PMPI_Bcast(buffer, count, datatype, root, comm);
}

int MPIAPI MPI_Scatter (void *sendbuf, int sendcnt, MPI_Datatype sendtype, 
                void *recvbuf, int recvcount, MPI_Datatype recvtype, 
                int root, MPI_Comm comm)
{
    if (Break("MPI_Scatter")) BREAK_HERE();
    return PMPI_Scatter(sendbuf, sendcnt, sendtype, recvbuf, recvcount, recvtype,
               root, comm);
}

int MPIAPI MPI_Scatterv (void *sendbuf, int *sendcnt, int *displs, MPI_Datatype sendtype, 
                void *recvbuf, int recvcount, MPI_Datatype recvtype, 
                int root, MPI_Comm comm)
{
    if (Break("MPI_Scatterv")) BREAK_HERE();
    return PMPI_Scatterv(sendbuf, sendcnt, displs, sendtype, recvbuf, recvcount,
               recvtype, root, comm);
}

int MPIAPI MPI_Gatherv (void *sendbuf, int sendcnt, MPI_Datatype sendtype, 
                void *recvbuf, int *recvcount, int *displs, MPI_Datatype recvtype, 
                int root, MPI_Comm comm)
{
    if (Break("MPI_Gatherv")) BREAK_HERE();
    return PMPI_Gatherv(sendbuf, sendcnt, sendtype, recvbuf, recvcount, displs, recvtype,
               root, comm);
}

int MPIAPI MPI_Gather (void *sendbuf, int sendcnt, MPI_Datatype sendtype, 
                void *recvbuf, int recvcount, MPI_Datatype recvtype, 
                int root, MPI_Comm comm)
{
    if (Break("MPI_Gather")) BREAK_HERE();
    return PMPI_Gather(sendbuf, sendcnt, sendtype, recvbuf, recvcount, recvtype,
               root, comm);
}

int MPIAPI MPI_Allgather (void *sendbuf, int sendcnt, MPI_Datatype sendtype, 
                void *recvbuf, int recvcount, MPI_Datatype recvtype, 
                MPI_Comm comm)
{
    if (Break("MPI_Allgather")) BREAK_HERE();
    return PMPI_Allgather(sendbuf, sendcnt, sendtype, recvbuf, recvcount,
               recvtype, comm);
}

int MPIAPI MPI_Allgatherv (void *sendbuf, int sendcnt, MPI_Datatype sendtype, 
                           void *recvbuf, int *recvcount, int *displs,
                           MPI_Datatype recvtype, MPI_Comm comm)
{
    if (Break("MPI_Allgatherv")) BREAK_HERE();
    return PMPI_Allgatherv(sendbuf, sendcnt, sendtype, recvbuf, recvcount, displs,
               recvtype, comm);
}

int MPIAPI MPI_Alltoall (void *sendbuf, int sendcnt, MPI_Datatype sendtype,
                         void *recvbuf, int recvcount, MPI_Datatype recvtype,
                         MPI_Comm comm)
{
    if (Break("MPI_Alltoall")) BREAK_HERE();
    return PMPI_Alltoall(sendbuf, sendcnt, sendtype, recvbuf, recvcount, recvtype,
               comm);
}

int MPIAPI MPI_Alltoallv (void *sendbuf, int *sendcnt, int *sdispls, MPI_Datatype sendtype,
                         void *recvbuf, int *recvcount, int *rdispls, MPI_Datatype recvtype,
                         MPI_Comm comm)
{
    if (Break("MPI_Alltoallv")) BREAK_HERE();
    return PMPI_Alltoallv(sendbuf, sendcnt, sdispls, sendtype, recvbuf, recvcount, rdispls,
               recvtype, comm);
}

int MPIAPI MPI_Scan (void *sendbuff, void *recvbuff, int count, 
    MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)
{
    if (Break("MPI_Scan")) BREAK_HERE();
    return PMPI_Scan(sendbuff, recvbuff, count, datatype, op, comm);
}

int MPIAPI MPI_Reduce (void *sendbuff, void *recvbuff, int count, 
    MPI_Datatype datatype, MPI_Op op, int root, MPI_Comm comm)
{
    if (Break("MPI_Reduce")) BREAK_HERE();
    return PMPI_Reduce(sendbuff, recvbuff, count, datatype, op, root, comm);
}

int MPIAPI MPI_Reduce_scatter (void *sendbuff, void *recvbuff, int *recvcnt, 
    MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)
{
    if (Break("MPI_Reduce_scatter")) BREAK_HERE();
    return PMPI_Reduce_scatter(sendbuff, recvbuff, recvcnt, datatype, op, comm);
}

int MPIAPI MPI_Allreduce (void *sendbuff, void *recvbuff, int count, 
    MPI_Datatype datatype, MPI_Op op, MPI_Comm comm)
{
    if (Break("MPI_Allreduce")) BREAK_HERE();
    return PMPI_Allreduce(sendbuff, recvbuff, count, datatype, op, comm);
}

int MPIAPI MPI_Waitall (int count, MPI_Request *array_of_requests, 
                            MPI_Status *array_of_statuses)
{
    if (Break("MPI_Waitall")) BREAK_HERE();
    return PMPI_Waitall(count, array_of_requests, array_of_statuses);
}

int MPIAPI MPI_Waitany (int count, MPI_Request *array_of_requests, 
                int *index, MPI_Status *status)
{
    if (Break("MPI_Waitany")) BREAK_HERE();
    return PMPI_Waitany(count, array_of_requests, index, status);
}

int MPIAPI MPI_Waitsome (int incount, MPI_Request *array_of_requests, 
                int *outcount, int *array_of_indices, MPI_Status *array_of_statuses)
{
    if (Break("MPI_Waitsome")) BREAK_HERE();
    return PMPI_Waitsome(incount, array_of_requests, outcount, array_of_indices,
               array_of_statuses);
}

int MPIAPI MPI_Testany (int count, MPI_Request *array_of_requests, 
        int *index, int *flag, MPI_Status *array_of_statuses)
{
    if (Break("MPI_Testany")) BREAK_HERE();
    return PMPI_Testany(count, array_of_requests, index, flag,
               array_of_statuses);
}

int MPIAPI MPI_Testsome (int incount, MPI_Request *array_of_requests, 
        int *outcount, int *array_of_indices, MPI_Status *array_of_statuses)
{
    if (Break("MPI_Testsome")) BREAK_HERE();
    return PMPI_Testsome(incount, array_of_requests, outcount, array_of_indices,
               array_of_statuses);
}

int MPIAPI MPI_Testall (int count, MPI_Request *array_of_requests, int *flag,
                            MPI_Status *array_of_statuses)
{
    if (Break("MPI_Testall")) BREAK_HERE();
    return PMPI_Testall(count, array_of_requests, flag, array_of_statuses);
}

int MPIAPI MPI_Abort (MPI_Comm comm, int errorcode)
{
    if (Break("MPI_Abort")) BREAK_HERE();
    return PMPI_Abort(comm, errorcode);
}

int MPIAPI MPI_Finalize()
{
    if (Break("MPI_Finalize")) BREAK_HERE();
    Disconnect(fd);
    return PMPI_Finalize();
}
